#include "common.hpp"

int task(size_t device_id, size_t data_size, size_t stream_num) {
    ACLCHECK(aclrtSetDevice(device_id));
    std::this_thread::sleep_for(std::chrono::seconds(3));
    auto start = std::chrono::high_resolution_clock::now();

    std::vector<aclrtStream> streams(stream_num);
    std::vector<std::thread> threads(stream_num);
    for (int i = 0; i < stream_num; i++) {
        ACLCHECK(aclrtCreateStream(&streams[i]));
    }

    // 申请Host内存, 并初始化, 让每一块的数据都不一样
    float *hostBuff;
    float *varifyBuff;
    float *deviceBuff;
    ACLCHECK(aclrtMallocHost((void **)&hostBuff, data_size * stream_num));
    ACLCHECK(aclrtMallocHost((void **)&varifyBuff, data_size * stream_num));
    ACLCHECK(aclrtMalloc((void **)&deviceBuff, data_size * stream_num, ACL_MEM_MALLOC_HUGE_FIRST));
    for (size_t i = 0; i < stream_num; i++) {
        auto func = [i, data_size, hostBuff, varifyBuff]() {
            for (size_t j = 0; j < data_size / sizeof(float); j++) {
                auto index = data_size * i / sizeof(float) + j;
                hostBuff[index] = varifyBuff[index] = (1E3 * i + j * 1E-6);
            }
        };
        threads[i] = std::thread(func);
    }
    for (int i = 0; i < stream_num; i++) {
        if (threads[i].joinable()) threads[i].join();
    }

    // 开始h2d的计时
    auto h2d_start = std::chrono::high_resolution_clock::now();

    // 将Host数据复制到Device侧
    for (int i = 0; i < stream_num; i++) {
        auto offset = data_size * i / sizeof(float);
        ACLCHECK(aclrtMemcpyAsync(deviceBuff + offset, data_size, hostBuff + offset, data_size, ACL_MEMCPY_HOST_TO_DEVICE, streams[i]));
    }
    for (int i = 0; i < stream_num; i++) {
        ACLCHECK(aclrtSynchronizeStream(streams[i]));
    }

    // 计算h2d的带宽
    std::chrono::duration<double, std::milli> h2d_dur;
    h2d_dur = std::chrono::high_resolution_clock::now() - h2d_start;
    float h2d_time = h2d_dur.count();

    // 清理Host内存
    for (int i = 0; i < stream_num; i++) {
        auto offset = data_size * i / sizeof(float);
        ACLCHECK(aclrtMemset(hostBuff + offset, data_size, 0x00, data_size));
    }

    // 开始d2h的计时
    auto d2h_start = std::chrono::high_resolution_clock::now();

    // 将Device侧数据复制回Host
    for (int i = 0; i < stream_num; i++) {
        auto offset = data_size * i / sizeof(float);
        ACLCHECK(aclrtMemcpyAsync(hostBuff + offset, data_size, deviceBuff + offset, data_size, ACL_MEMCPY_DEVICE_TO_HOST, streams[i]));
    }
    for (int i = 0; i < stream_num; i++) {
        ACLCHECK(aclrtSynchronizeStream(streams[i]));
    }

    // 计算d2h的带宽
    std::chrono::duration<double, std::milli> d2h_dur;
    d2h_dur = std::chrono::high_resolution_clock::now() - d2h_start;
    float d2h_time = d2h_dur.count();

    // 校验Host侧的数据与预期是否一致
    for (int i = 0; i < stream_num; i++) {
        auto func = [i, data_size, hostBuff, varifyBuff]() {
            for (uint32_t j = 0; j < data_size / sizeof(float); j++) {
                auto index = i * data_size / sizeof(float) + j;
                if (hostBuff[index] - varifyBuff[index] > 1E-12) {
                    printf("error: i=%d j=%d host_value=%f varify_value=%f\n", i, j, hostBuff[index], varifyBuff[index]);
                    break;
                }
            }
        };
        threads[i] = std::thread(func);
    }
    for (int i = 0; i < stream_num; i++) {
        if (threads[i].joinable()) threads[i].join();
    }

    // 清理数据
    ACLCHECK(aclrtFree(deviceBuff));
    ACLCHECK(aclrtFreeHost(hostBuff));
    ACLCHECK(aclrtFreeHost(varifyBuff));
    for (int i = 0; i < stream_num; i++) {
        ACLCHECK(aclrtDestroyStream(streams[i]));
    }
    ACLCHECK(aclrtResetDevice(device_id));

    // 打印测试结果
    float total_size = stream_num * (float)data_size / GB;
    float total_time = (std::chrono::high_resolution_clock::now() - start).count() * 1E-6;
    printf("device: %d, stream_num: %d, data_size(MB): %.2f, total_time(ms): %.3f, h2d_time(ms): %.3f, d2h_time(ms): %.3f, h2d_bandwidth(GB/s): %.3f, d2h_bandwidth(GB/s): %.3f\n",
           device_id, stream_num, (float)data_size / MB, total_time, h2d_time, d2h_time, total_size / (h2d_time * 1E-3), total_size / (d2h_time * 1E-3));

    return 0;
}

int main(int argc, char *argv[]) {
    if (argc != 4) {
        printf("./bin/main <device_id> <data_size_MB> <stream_num>\n");
        return 1;
    }
    size_t device_id = std::stoi(std::string(argv[1]));                 // 获取设备的ID
    size_t data_size = (size_t)(std::stoi(std::string(argv[2]))) * MB;  // 测试数据大小, Byte
    size_t stream_num = std::stoi(std::string(argv[3]));                // 修改该值以并行启动多条流

    task(device_id, data_size, stream_num);
    return 0;
}
